
#include "shading.fx"
#include "fog.fx"

static const unsigned int LEAF_GENERATOR = 0;
static const unsigned int LEAF_MAIN		 = 1;
static const unsigned int LEAF_SIDE      = 2;

static const float y_min =		2.95f;
static const float y_range =	1.0f;
static const float r =			0.15f;
static const int   N_LEAVES =	32;


float arrLeavesHeights[N_LEAVES];


struct Leaf
{
	float4	pos		: POSITION;
	uint	type	: TYPE;
	float  length : LENGTH;
	float  interval : INTERVAL;
	uint	id		: ID;
};

struct GS_OUT
{
	float3 position : POSITION;
	float3 normal	: NORMAL;
	float3 tangent : TANGENT;
	float2 tex		: TEX_COORDS;
};

struct VS_IN
{
	float3 position : POSITION;
	float3 normal	: NORMAL;
	float3 tangent : TANGENT;
	float2 tex		: TEX_COORDS;
};

struct VS_OUT
{
	float4 position : SV_POSITION;
	float4 ws_pos	: WS_POSITION;
	float3 normal	: NORMAL;
	float2 tex		: TEX_COORDS;
	uint   id		: ID;
};


float4		p_eye_pos;
float4x4	m_VP; 
Texture2D	noise_texture;
float4x4	m_WVP;
float		dt;
float		grow_speed;
float4x4	mtxLeaves[N_LEAVES];





///////////////////////

float4x4 m_World;
float4x4 m_WorldViewProj;




SamplerState s_volume
{
    Filter = MIN_MAG_MIP_POINT;
    AddressU = Clamp;
    AddressV = Clamp;
    AddressW = Clamp;
};




DepthStencilState disabled
{
	DepthEnable = false;
	DepthWriteMask = false;
};






///==============================================================================================================================
///---------------- D   R   A   W -----------------------------------------------------------------------------------------------



VS_OUT VS_Draw(VS_IN vin)
{
	VS_OUT vout;

	float id = vin.tangent.x;


	vout.position = mul(float4(vin.position, 1.0f), mtxLeaves[id]);
	vout.position = mul(vout.position,  m_WorldViewProj);
	
	vout.ws_pos = mul(float4(vin.position, 1.0f), m_World);
	vout.normal = mul(float4(vin.normal, 0.0f), m_World).xyz;
	vout.tex =	  vin.tex;
	vout.id = id;

	return vout;
}

float4 Leaves_PS(VS_OUT pin) : SV_TARGET
{
	// color = shade() : diffuse+specular
	float4 color = float4(shade(pin.ws_pos.xyz, normalize(pin.normal), 0), 1.0f);
			
	
	float4 ambient = float4(0.1f, 0.1f, 0.1f, 0.0f);

	// color	= color	   * tex_diffuse
	// ambient	= ambient  * tex_diffuse
	if (mat_params.vIsTextured.x)
	{
		float4 tex_diff = tex_diffuse.Sample(s_linear, pin.tex);
		color *= tex_diff;
		ambient *= tex_diff;
	}

	// set alpha in w
	color = float4(color.xyz, mat_params.diffuse.w);

	// color = color + ambient
	color += ambient;

	return color;
}


RasterizerState rs_Wireframe
{
	CullMode = Back;
	FillMode = Wireframe;
};


technique10 Leaves_Draw
{
	pass P0
	{
		SetVertexShader(CompileShader( vs_4_0, VS_Draw() ) );
		SetGeometryShader( NULL );
		SetPixelShader( CompileShader( ps_4_0, Leaves_PS() ) );

		SetDepthStencilState(0, 0);
	}
}


///==============================================================================================================================
///---------------- LEAVES   G E O M E T R Y   S O -----------------------------------------------------------------------------------------------

struct Leaf_VS
{
	float4	pos		: POSITION;
	uint	type	: TYPE;
	float  length : LENGTH;
	float  interval : INTERVAL;
	uint	id		: ID;
};

struct Leaf_GS
{
	float4	pos		: POSITION;
	uint	type	: TYPE;
	float  length : LENGTH;
	float  interval : INTERVAL;
	uint	id		: ID;
};

Leaf_GS Leaves_VS(Leaf_VS vin)
{
	Leaf_GS vout;

	vout.pos = vin.pos;
	vout.type = vin.type;
	vout.interval = vin.interval;
	vout.length = vin.length;
	vout.id = vin.id;
	return vout;
}


[maxvertexcount(48)]
void GS(point Leaf_GS gin[1],	inout PointStream<GS_OUT> gsout)
{
	
	// input leaf element
	Leaf leaf;
	leaf.interval = gin[0].interval;
	leaf.length = gin[0].length;
	leaf.pos = gin[0].pos;
	leaf.type = gin[0].type;

	float3 x = float3(1.0f, 0.0f, 0.0f);
	float3 z = float3(0.0f, 0.0f, 1.0f);

	float size = 0.03f;

	// only side leaves are transformed to geometry
	if (leaf.type == LEAF_SIDE)
	{
		// POS.xyzw     : main leaf position + length
		// INTERVAL     : maximum side leaf length
		// LENGTH       : length

		// size = max length
		size = leaf.interval;
		
		float max_slope = 1.0f;
		float min_slope = 0.01f;

		// y of the main leaf in [0,1] range (depending on height min,max)
		float height_norm = (leaf.pos.y - y_min) / y_range; // [0,1]
		
		// compute alpha for -(a*x)^2 which describes the main leaf shape
		float alpha = max_slope - (max_slope - min_slope) * height_norm; // [max_slope, min_slope]
		
		////
		float3 p0 = gin[0].pos.xyz;
		
		// h = -(a*l)^2
		float h = -alpha * gin[0].pos.w * gin[0].pos.w;
		
		float3 dir = normalize(float3(gin[0].pos.x, 0, gin[0].pos.z));

		leaf.pos.xyz = p0 + float3(dir.x * gin[0].pos.w, h, dir.z * gin[0].pos.w);
		
		// in local pos - will be saved in tangent var 
		float cntL = 0.0f;
		float3 cntPoint = p0; 
		leaf.pos.xz = leaf.pos.xz - cntPoint.xz;

		///////

		// f(x) = -(ax)^2 => f'(x) = -2*a^2 * x
		float dL = -2.0f * alpha * alpha * leaf.pos.w;

		float3 tangent = normalize(float3(dir.x, dL, dir.z));
		float3 right = cross(float3(0,1,0), tangent);
		float3 up = cross(tangent, right);

		/////////////////	

		/////////////////

		float l_size = 0.09f; 

		float3 verts[6];
		float3 v_right = leaf.pos.xyz + right * leaf.length * 1.5f;		//	
		float3 v_left  = leaf.pos.xyz - right * leaf.length * 1.5f;		//
														
		float3 v_up    = leaf.pos.xyz + up.xyz * 0.01f;			
		float3 v_down  = leaf.pos.xyz - up.xyz * 0.01f;			
														
		float3 v_front = leaf.pos.xyz + tangent * l_size;			//			
		float3 v_back  = leaf.pos.xyz - tangent * l_size;			//


		float4 tri[24];	float2 uv[24];
		
		tri[0] = float4(v_up, 1.0f);		uv[0] = float2(0.5f, 0.5f);
		tri[1] = float4(v_front, 1.0f);		uv[1] = float2(0.0f, 0.5f);
		tri[2] = float4(v_right, 1.0f);		uv[2] = float2(0.5f, 0.0f);

		tri[3] = float4(v_up, 1.0f);		uv[3] = float2(0.5f, 0.5f);
		tri[4] = float4(v_back, 1.0f);		uv[4] = float2(1.0f, 0.5f);
		tri[5] = float4(v_right, 1.0f);		uv[5] = float2(0.5f, 0.0f);

		tri[6] = float4(v_up, 1.0f);		uv[6] = float2(0.5f, 0.5f);
		tri[7] = float4(v_front, 1.0f);		uv[7] = float2(0.0f, 0.5f);
		tri[8] = float4(v_left, 1.0f);		uv[8] = float2(0.5f, 1.0f);

		tri[9] =  float4(v_up, 1.0f);		uv[9] = float2(0.5f, 0.5f);
		tri[10] = float4(v_back, 1.0f);		uv[10] = float2(1.0f, 0.5f);
		tri[11] = float4(v_left, 1.0f);		uv[11] = float2(0.5f, 1.0f);

		///////////////

		tri[12] = float4(v_down, 1.0f);		uv[12] = float2(0.5f, 0.5f);
		tri[13] = float4(v_front, 1.0f);	uv[13] = float2(0.0f, 0.5f);
		tri[14] = float4(v_right, 1.0f);	uv[14] = float2(0.5f, 0.0f);

		tri[15] = float4(v_down, 1.0f);		uv[15] = float2(0.5f, 0.5f);
		tri[16] = float4(v_back, 1.0f);		uv[16] = float2(1.0f, 0.5f);
		tri[17] = float4(v_right, 1.0f);	uv[17] = float2(0.5f, 0.0f);

		tri[18] = float4(v_down, 1.0f);		uv[18] = float2(0.5f, 0.5f);
		tri[19] = float4(v_front, 1.0f);	uv[19] = float2(0.0f, 0.5f);
		tri[20] = float4(v_left, 1.0f);		uv[20] = float2(0.5f, 1.0f);

		tri[21] =  float4(v_down, 1.0f);	uv[21] = float2(0.5f, 0.5f);
		tri[22] = float4(v_back, 1.0f);		uv[22] = float2(1.0f, 0.5f);
		tri[23] = float4(v_left, 1.0f);		uv[23] = float2(0.5f, 1.0f);


		float id = gin[0].id;

		for (int i=0; i<24; i++)
		{
			GS_OUT gout;
			gout.position = tri[i].xyz;
			gout.normal = up;
			gout.tangent = float3(id,id,id); // tangent is used to save the id, instead of using multiple SO streams
													// which is DX11 only feature
			gout.tex = uv[i];
			gsout.Append(gout);
		}

	}
}

GeometryShader gsStreamOutDraw = ConstructGSWithSO( 
				CompileShader( gs_4_0, GS() ), 
				"POSITION.xyz; NORMAL.xyz; 0:TANGENT.xyz; TEX_COORDS.xy"
				);



technique10 Leaves_SO_Geometry
{
	pass P0
	{
		SetVertexShader( CompileShader ( vs_4_0, Leaves_VS() ) );
		SetGeometryShader( gsStreamOutDraw );
		SetPixelShader ( NULL );

		SetDepthStencilState(disabled, 0);
	}

}



///==============================================================================================================================
///---------------- LEAVES     SO       U   P  D   A   T   E -----------------------------------------------------------------------------------------------


[maxvertexcount(50)]
void GS_SO(point Leaf gin[1], inout PointStream<Leaf> gsout)
{
	Leaf leaf;
	leaf = gin[0];
	

	if (gin[0].type == LEAF_GENERATOR)
	{
		// create N main leaves randomly placed - set length to a predefined value so new leaves are generated only once
		int n_leaves = N_LEAVES;
		
		float angle_step = 2.0f*3.1415f / (float)n_leaves;
		
		if (leaf.length < 0.1f)
		{
			// create main leaves
			for (int i=0; i<n_leaves; i++)
			{
				Leaf main_leaf;

				main_leaf.type = LEAF_MAIN;
				main_leaf.length = 0.0f;
				main_leaf.interval = 0.0f;
				
				// get random variable theta : 0 <= theta <= 360 
				float theta = i * angle_step;

				// get random variable z	 : min_leaf <= z <= max_leaf
				float rnd_y = noise_texture.SampleLevel(s_linear, float2(theta, -theta), 0).r;
				float y_value = y_min + rnd_y * y_range;
				
				y_value = arrLeavesHeights[i];

				main_leaf.pos = float4(r*cos(theta), /*3.0*/y_value, r*sin(theta), 1.0f);
				
				main_leaf.id = i;

				gsout.Append(main_leaf);
			}
			
		}
		leaf.length = 1.0f; // do not create new leaves in each frame
		gsout.Append(leaf);	

	}
	else if (gin[0].type == LEAF_MAIN)
	{
		
		if (leaf.length <= 1.0f)
		{
			// increase length
			leaf.length += dt * grow_speed;
			// increase interval
			leaf.interval += dt * grow_speed;

			// if length > threshold create new side leaf
			if (leaf.interval >= 0.01f)
			{
				int new_side_leaves = (int)(leaf.interval / 0.01f);
				
				for (int i=0; i<new_side_leaves; ++i)
				{

					// create 1 new leaf
					Leaf s_leaf;

					// TYPE
					s_leaf.type = LEAF_SIDE;
				
					float new_l = leaf.length - 0.01f * i;

					// POSITION : save parent's leaf pos + length so side leaf pos can be reconstructed 
					s_leaf.pos = float4(leaf.pos.xyz, new_l);

					// LENGTH (=0)
					s_leaf.length = 0.0f;
			
					float f_a = 0.4f;
					float f_b = 0.9f;
					float scale_factor = (1.57f + f_b) / f_a;

					// evaluating for main leaf's length gives maximum side length
					float max_l = cos(new_l*scale_factor*f_a - f_b);
					max_l *= 0.1f;
				
					// interval for side leaves holds the maximum length
					s_leaf.interval = max_l;
					s_leaf.id = leaf.id;

					gsout.Append(s_leaf);	
				}
				// reset interval
				leaf.interval = 0.0f;
			}
		}
		gsout.Append(leaf);

	}
	else if (gin[0].type == LEAF_SIDE)
	{
		// increase length
		if (leaf.length < leaf.interval) {
			leaf.length += dt * grow_speed;
			leaf.length = min(leaf.length, leaf.interval);
		}

		gsout.Append(leaf);
	}
}



GeometryShader gsStreamOut = ConstructGSWithSO( 
	CompileShader( gs_4_0, GS_SO() ), 
	"POSITION.xyzw; TYPE.x; LENGTH.x; INTERVAL.x; ID.x");





technique10 Leaves_SO
{
	pass P0
	{
		SetVertexShader( CompileShader ( vs_4_0, Leaves_VS() ) );
		SetGeometryShader( gsStreamOut );
		SetPixelShader ( NULL);

		SetDepthStencilState(disabled, 0);
	}
}
